# app/nx_planner/map_config.py

import json

class MapConfig:
    def __init__(self, FILE_NAME: str="app/nx_planner/data/map_config_tj.json"):
        self.file_name = FILE_NAME
        self.map_config = json.load(open(self.file_name))
        self.nodes = self.map_config["nodes"]
        self.edges = self.map_config["edges"]

    def save_map_config(self, CONNECT_NUMS: int = 8) -> None:
        """
        保存地图配置

        :params CONNECT_NUMS: 通道连接数
        """
        print(f"地图节点：{self.map_config['nodes']}")
        print(f"原地图边：{self.map_config['edges']}")

        if CONNECT_NUMS == 4:
            new_map_config= self.complete_grid_edges_4_connected(self.map_config)
        elif CONNECT_NUMS == 8:
            new_map_config = self.complete_grid_edges_8_connected(self.map_config)
        else:
            raise ValueError("Invalid CONNECT_NUMS")

        print(f"地图节点：{new_map_config['nodes']}")
        print(f"新增边：{new_map_config['edges']}")

        json.dump(new_map_config, open(self.file_name, "w"))


    def complete_grid_edges_4_connected(self, MAP_CONFIG: dict) -> dict:
        """
        自动补全网格图的边（四连通）
        
        Params:
            MAP_CONFIG: 包含"nodes"和"edges"键的字典
            
        Returns:
            补全后的完整map_config字典
        """
        map_config = MAP_CONFIG
        nodes = map_config["nodes"]
        existing_edges = set(tuple(edge) for edge in map_config["edges"]) # 转换为元组并去重
        
        # 创建一个字典来快速查找节点是否存在（虽然你的节点列表似乎是顺序的，但这样做更稳健）
        node_set = set(nodes)
        
        # 从节点列表中解析出行列范围
        rows = sorted(set(int(node.split(',')[0]) for node in nodes))
        cols = sorted(set(int(node.split(',')[1]) for node in nodes))
        max_row = max(rows)
        max_col = max(cols)
        
        new_edges = []
        # 也可以直接遍历nodes列表，这里选择遍历行列更直观
        for r in rows:
            for c in cols:
                current_node = f"{r},{c}"
                if current_node not in node_set:
                    continue # 跳过不存在的节点（虽然你的列表是完整的）
                    
                # 1. 检查并添加向右的边 (连接到右邻居: r, c+1)
                if c < max_col:
                    right_node = f"{r},{c+1}"
                    if right_node in node_set:
                        edge_tuple_forward = (current_node, right_node)
                        edge_tuple_backward = (right_node, current_node) # 无向图通常不考虑方向，但确保唯一性
                        # 检查新边是否已存在（正反方向）
                        if edge_tuple_forward not in existing_edges and edge_tuple_backward not in existing_edges:
                            new_edges.append(list(edge_tuple_forward)) # 以列表形式添加
                            existing_edges.add(edge_tuple_forward) # 添加到已有边集合，避免后续重复检查
                            
                # 2. 检查并添加向下的边 (连接到下邻居: r+1, c)
                if r < max_row:
                    down_node = f"{r+1},{c}"
                    if down_node in node_set:
                        edge_tuple_forward = (current_node, down_node)
                        edge_tuple_backward = (down_node, current_node)
                        if edge_tuple_forward not in existing_edges and edge_tuple_backward not in existing_edges:
                            new_edges.append(list(edge_tuple_forward))
                            existing_edges.add(edge_tuple_forward)
                
        print(f"新增了 {len(new_edges)} 条边")

        # 将新边添加到原始的edges列表中
        map_config["edges"].extend(new_edges)
        
        return map_config


    def complete_grid_edges_8_connected(self, MAP_CONFIG: dict) -> dict:
        """
        自动补全网格图的边（八连通）
        
        Params:
            MAP_CONFIG: 包含"nodes"和"edges"键的字典
            
        Returns:
            补全后的完整map_config字典
        """
        map_config = MAP_CONFIG
        nodes = map_config["nodes"]
        # 将现有边转换为集合以便快速查找（考虑无向边的两种方向）
        existing_edges_set = set()
        for edge in map_config["edges"]:
            # 将边转换为元组并排序，确保 (a, b) 和 (b, a) 被视为同一条边
            sorted_edge = tuple(sorted(edge))
            existing_edges_set.add(sorted_edge)
        
        # 创建节点集合用于快速查找
        node_set = set(nodes)
        
        # 解析出行列范围
        rows = sorted(set(int(node.split(',')[0]) for node in nodes))
        cols = sorted(set(int(node.split(',')[1]) for node in nodes))
        max_row = max(rows)
        max_col = max(cols)
        
        new_edges = []
        
        # 八连通的方向偏移量: (行偏移, 列偏移)
        # 包括: 右, 右下, 下, 左下, 左, 左上, 上, 右上
        directions = [
            (0, 1),   # 右
            (1, 1),   # 右下
            (1, 0),   # 下
            (1, -1),  # 左下
            (0, -1),  # 左
            (-1, -1), # 左上
            (-1, 0),  # 上
            (-1, 1)   # 右上
        ]
        
        # 遍历每个节点
        for r in rows:
            for c in cols:
                current_node = f"{r},{c}"
                if current_node not in node_set:
                    continue
                    
                # 检查所有8个方向
                for dr, dc in directions:
                    new_r, new_c = r + dr, c + dc
                    
                    # 检查新坐标是否在网格范围内
                    if new_r < min(rows) or new_r > max_row or new_c < min(cols) or new_c > max_col:
                        continue
                        
                    neighbor_node = f"{new_r},{new_c}"
                    
                    # 检查邻居节点是否存在
                    if neighbor_node not in node_set:
                        continue
                        
                    # 创建边的两种可能表示（无向图）
                    edge_tuple = tuple(sorted([current_node, neighbor_node]))
                    
                    # 检查边是否已存在
                    if edge_tuple not in existing_edges_set:
                        new_edges.append([current_node, neighbor_node])
                        existing_edges_set.add(edge_tuple)  # 添加到已有边集合，避免重复添加
        
        # 将新边添加到原始的edges列表中
        map_config["edges"].extend(new_edges)
        
        return map_config